import matplotlib.pyplot as plt
import numpy as np

# 基函数  de Boor-Cox 递推定义
def b_basis(i, p, u, U):
    # U 为[u0,u1,u2,...,u_{n+p+1}]
    if p == 0:
        basis_i = 1 if U[i]<u<=U[i+1] else 0
    else:
        # 计算支撑区间长度
        length1 = U[i+p] - U[i]
        length2 = U[i+p+1] - U[i+1]
        # 定义基函数中的两个系数
        alpha = 0 if not length1 else (u-U[i])/length1 
        beta = 0 if not length2 else (U[i+p+1]-u)/length2
        # 递归定义
        basis_i = alpha * b_basis(i, p-1, u, U) + beta * b_basis(i+1, p-1, u, U)
    return basis_i

# B样条曲线
def b_spline(U, Pts):
    '''
        U  : list(u0,u1,u2,...,u_{n+p+1}),len = p+n+2
        Pts: list((xi,yi,zi)) 控制点
    '''
    p_num = len(Pts) # n+1
    p = len(U)-p_num-1 # 阶数
    u_num = 100 # 插值数量
    points = np.zeros((u_num,3))  # 存放B样条的坐标
    Bi = np.zeros((u_num,1))
    for i in range(p_num):
        j = 0
        # u 取值
        for u in np.linspace(U[0], U[-1], 100):
            Bi[j,0] = b_basis(i, p, u, U)
            j = j + 1
        points += Bi @ [Pts[i]]
    return points


if __name__ == '__main__':
    # 创建画布
    plt.rcParams["font.sans-serif"]=["SimHei"]  # 设置字体
    plt.rcParams["axes.unicode_minus"]=False    # “-”负号的显示
    fig = plt.figure()                          # 空的画布
    ax = fig.add_subplot(111, projection='3d')  # 创建3d绘图区域
    ax.tick_params(labelsize=8)                 # 刻度字体

    # 绘图
    U = [0, 0, 0, 0, 1, 2, 3, 4, 5, 5, 5, 5]
    Pts = np.array([(0,0,0),(1,2,1),(5,2,2),(5,4,3),
                    (8,4,4),(4,5,5),(3,5,6),(3,9,7)])
    points = b_spline(U, Pts)
    ax.plot(Pts[:,0],Pts[:,1],Pts[:,2])
    ax.plot(points[:,0],points[:,1],points[:,2])
    ax.set_title('B 样条曲线')

    # 参考线
    def axs(xx=(-1,1),yx=(-1,1),zx=(-1,1)): 
        color = ['r', 'g', 'b']
        for i,t in  enumerate((xx, yx, zx)):
            xs = [(0,0), (0,0), (0,0)]
            f = True if t[0]<0 and t[1]>0 else False
            if f:
                xs[i] = (t[0], 0)
                ax.plot(xs[0], xs[1], xs[2], linestyle=':', color=color[i])
                xs[i] = (0, t[1])
                ax.plot(xs[0], xs[1], xs[2], linestyle='-', color=color[i])
            else: 
                xs[i] = (t[0], t[1])
                ax.plot(xs[0], xs[1], xs[2], linestyle='-', color=color[i])
    axs((0,5),(0,5),zx=(0,5))
    ax.grid(b=None) # 隐藏网格

    # 显示窗口
    plt.show()

